package cmd

import (
	"bytes"
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/spf13/cobra"
	"github.com/tikv/client-go/txnkv"
)

var backupCmd = &cobra.Command{
	Use:   "backup",
	Short: "Backup He3DB",
	Long:  "Welcome to use hbr for He3DB backup",
	Run:   runBackup,
}

func init() {

	rootCmd.AddCommand(backupCmd)
}

func runBackup(cmd *cobra.Command, args []string) {
	access_key, _ := cmd.Flags().GetString("access_key")
	secret_key, _ := cmd.Flags().GetString("secret_key")
	endpoint, _ := cmd.Flags().GetString("endpoint")
	region, _ := cmd.Flags().GetString("region")
	bucket, _ := cmd.Flags().GetString("bucket")
	pd, _ := cmd.Flags().GetString("pd")
	backup_name, _ := cmd.Flags().GetString("name")

	if access_key == "" || secret_key == "" || endpoint == "" || region == "" || bucket == "" || pd == "" || backup_name == "" {
		fmt.Printf("PARAMETER ERROR!\n")
		return
	}

	client, err := txnkv.NewClient([]string{pd})
	if err != nil {
		fmt.Printf("Connect Tikv Error!\n%v\n", err)
		return
	}
	defer func() {
		client.Close()
	}()

	sess, err := session.NewSession(&aws.Config{
		Region:      aws.String(region),
		Endpoint:    aws.String(endpoint),
		Credentials: credentials.NewStaticCredentials(access_key, secret_key, ""),
		S3ForcePathStyle: aws.Bool(true),
	})
	if err != nil {
		fmt.Printf("Connect S3 Error!\n%v\n", err)
		return
	}
	s3_client := s3.New(sess)

	ts, err := client.CurrentTimestamp("global")
	if err != nil {
		fmt.Printf("Set Timestamp Error!\n%v\n", err)
		return
	}
	snapshot := client.GetSnapshot(ts)

	iter, err := snapshot.Iter([]byte(""), []byte(""))
	if err != nil {
		fmt.Printf("Iter Error!\n%v\n", err)
		return
	}
	defer iter.Close()
	count := 0
	for iter.Valid() {
		k, v := iter.Key(), iter.Value()
		var filename string = ""

		filename = fmt.Sprintf("%x", k)
		fmt.Printf("filename:%s\n", filename)

		_, err = s3_client.PutObject(&s3.PutObjectInput{
			Bucket: aws.String(bucket), // bucket名称
			Key:    aws.String(backup_name + "/" + filename),
			Body:   bytes.NewReader(v),
		})
		if err != nil {
			fmt.Printf("S3 PutObject Error!\n%v\n", err)
			return
		}

		if err := iter.Next(); err != nil {
			fmt.Printf("Iter Next Error!\n%v\n", err)
			return
		}
		count++
	}
	fmt.Printf("N:%v\n", count)
}
